import pickle
import numpy as np
import torch
import os

from Config.Config import ACTNET200V13_PKL,TRAIN_PROPOSALS_PKL,TEST_PROPOSALS_PKL,VAL_PROPOSALS_PKL,TSNscore_FEATURE_PATH
from DataSet.Tools_TSNscore import load_TSNscore_from_file,extract_TSN_feature
from DataSet.Utils import mulity_scale_pairs

class Dataset_TSNscore(object):
    '''
    取数据, 分为test和train两种模式
    '''
    def __init__(self,subset,modality):
        '''
        :param modality:  test or train
        :param subset: testing validation traning
        '''
        self.modality = modality
        self.subset = subset

        self.load_groundtruth()
        self.load_proposals()

    def load_groundtruth(self):
        with open(ACTNET200V13_PKL,'rb') as f:
            groundtruth = pickle.load(f)['database']
        # filter by subset
        remainkeyset = []
        for key in groundtruth.keys():
            if groundtruth[key]['subset'] == self.subset:
                remainkeyset.append(key)
        self.groundtruth = dict()
        for key in remainkeyset:
            self.groundtruth[key] = groundtruth[key]
        self.vids = remainkeyset

    def load_proposals(self):
        PKL = None
        if self.subset == 'training':
            PKL = TRAIN_PROPOSALS_PKL
        elif self.subset == 'testing':
            PKL = TEST_PROPOSALS_PKL
        elif self.subset == 'validation':
            PKL = VAL_PROPOSALS_PKL
        with open(PKL,'rb') as f:
            self.proposals = pickle.load(f)

    def nextbatch(self):
        pass

    def enum_genfucntion(self):
        pass
